package com.yjzx.llama.tensor.cuda;

import com.yjzx.llama.interfaces.AggregateFunction;
import com.yjzx.llama.interfaces.MapFunction;
import com.yjzx.llama.interfaces.MapWithIndexFunction;
import com.yjzx.llama.tensor.FloatTensor;
import com.yjzx.llama.util.ParallelUtil;
import jdk.incubator.vector.VectorShape;
import jdk.incubator.vector.VectorSpecies;

import java.util.Arrays;

/**
 * @author yjzx
 * @description BaseCudaFloatTensor
 * @since 2025/5/28
 */
public abstract class BaseCudaFloatTensor implements FloatTensor {
    /**
     * 获取llama.VectorBitSize的长度 vector bit大小
     */
    static final int VECTOR_BIT_SIZE = Integer.getInteger("llama.VectorBitSize", VectorShape.preferredShape().vectorBitSize());

    static final boolean USE_VECTOR_API = VECTOR_BIT_SIZE != 0;

    //static final ValueLayout.OfFloat JAVA_FLOAT_LE = ValueLayout.JAVA_FLOAT.withOrder（ByteOrder.LITTLE_ENDIAN）;
    //static final ValueLayout.OfShort JAVA_SHORT_LE = ValueLayout.JAVA_SHORT.withOrder（ByteOrder.LITTLE_ENDIAN）;

    // 快速乘法例程的首选向量大小.
    // （苹果芯片）NEON 仅支持最高 128 位的矢量.
    static final VectorSpecies<Float> F_SPECIES;
    static final VectorSpecies<Integer> I_SPECIES;
    static final VectorSpecies<Short> S_SPECIES_HALF;

    static {
        if (USE_VECTOR_API) {
            F_SPECIES = VectorShape.forBitSize(VECTOR_BIT_SIZE).withLanes(float.class);
            I_SPECIES = F_SPECIES.withLanes(int.class);
            S_SPECIES_HALF = VectorShape.forBitSize(F_SPECIES.vectorBitSize() / 2).withLanes(short.class);
            assert F_SPECIES.length() == S_SPECIES_HALF.length();
        } else {
            F_SPECIES = null;
            I_SPECIES = null;
            S_SPECIES_HALF = null;
        }
    }

    /**
     * 元素乘积
     *
     * @param dimensions 参数
     * @return 对所有参数进行乘积
     */
    public static int numberOfElements(int... dimensions) {
        assert Arrays.stream(dimensions).allMatch(i -> i > 0);
        return Arrays.stream(dimensions).parallel().reduce(Math::multiplyExact).orElseThrow();
    }

    /**
     * 点积量（默认的cpu求）
     *
     * @param thiz       张量1
     * @param thisOffset 张量1偏移起始量
     * @param that       张量2
     * @param thatOffset 张量2偏移起始量
     * @param size       数量
     * @return 对俩张量偏移后的数值进行求积后求和相加
     */
    public static float scalarDot(FloatTensor thiz, int thisOffset, FloatTensor that, int thatOffset, int size) {
        float result = 0f;
        for (int j = 0; j < size; j++) {
            result += thiz.getFloat(thisOffset + j) * that.getFloat(thatOffset + j);
        }
        return result;
    }

    /**
     * 点积函数
     *
     * @param thisOffset 张量1偏移起始量
     * @param that       张量2
     * @param thatOffset 张量2偏移起始量
     * @param size       数量
     * @return 对俩张量偏移后的数值进行求积后求和相加
     */
    @Override
    public float dot(int thisOffset, FloatTensor that, int thatOffset, int size) {
        return scalarDot(this, thisOffset, that, thatOffset, size);
    }

    /**
     * 矩阵乘积
     *
     * @param that 参数1
     * @param out  输出参数
     * @param dim0 0-dim0的每个值对
     * @param dim1
     */
    @Override
    public void matmul(FloatTensor that, FloatTensor out, int dim0, int dim1) {
        ParallelUtil.parallelFor(0, dim0, i -> out.setFloat(i, dot(i * dim1, that, 0, dim1)));
    }

    /**
     * 双重矩阵乘积
     *
     * @param context 0-dim0*context
     * @param that    参数1
     * @param out     输出参数
     * @param dim0    第一个
     * @param dim1    第二个
     */
    @Override
    public void matmul(int context, FloatTensor[] that, FloatTensor[] out, int dim0, int dim1) {
        if (that.length != out.length) {
            throw new IllegalArgumentException(String.format("that.len=%d, out.len=%d", that.length, out.length));
        }
        ParallelUtil.parallelForLong(0, (long) dim0 * context, ti -> {
            int idxArr = (int) (ti / dim0);
            int i = (int) (ti % dim0);
            out[idxArr].setFloat(i, dot(i * dim1, that[idxArr], 0, dim1));
        });
    }

    /**
     * 对
     *
     * @param thisOffset 当前对象偏移量
     * @param size       长度
     * @param seed       运算参数
     * @param reduce     进行的运算
     * @return 对当前对象指定的索引位置的值与seed在值size长度遍历进行reduce运算
     */
    @Override
    public float reduce(int thisOffset, int size, float seed, AggregateFunction reduce) {
        float result = seed;
        for (int i = 0; i < size; ++i) {
            result = reduce.apply(result, getFloat(thisOffset + i));
        }
        return result;
    }

    @Override

    public float sum(int thisOffset, int size) {
        return reduce(thisOffset, size, 0f, Float::sum);
    }

    @Override
    public float max(int thisOffset, int size) {
        return reduce(thisOffset, size, Float.NEGATIVE_INFINITY, Float::max);
    }

    /**
     * 对that的到that Offset到size长度偏移量的值替换为 this的偏移量及size长度偏移量的值
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     */
    @Override
    public void copyTo(int thisOffset, FloatTensor that, int thatOffset, int size) {
        that.mapWithIndexInPlace(thatOffset, size, (value, index) -> this.getFloat(index - thatOffset + thisOffset));
    }

    /**
     * @param thisOffset 偏移量
     * @param size       长度
     * @return 获取thisOffset到指定size长度中最大索引值
     */
    @Override
    public int argmax(int thisOffset, int size) {
        assert size > 0;
        int maxIndex = thisOffset;
        float maxValue = this.getFloat(maxIndex);
        int endIndex = thisOffset + size;
        for (int i = thisOffset; i < endIndex; ++i) {
            float f = this.getFloat(i);
            if (f > maxValue) {
                maxValue = f;
                maxIndex = i;
            }
        }
        return maxIndex;
    }

    /**
     * @return 获取有最大值的索引（最后一个）
     */
    @Override
    public int argmax() {
        return argmax(0, size());
    }

    @Override
    public FloatTensor mapInPlace(int thisOffset, int size, MapFunction mapFunction) {
        int endIndex = thisOffset + size;
        for (int i = thisOffset; i < endIndex; ++i) {
            setFloat(i, mapFunction.apply(getFloat(i)));
        }
        return this;
    }

    @Override
    public FloatTensor mapInPlace(MapFunction mapFunction) {
        return mapInPlace(0, size(), mapFunction);
    }

    @Override
    public FloatTensor mapWithIndexInPlace(int thisOffset, int size, MapWithIndexFunction mapWithIndexFunction) {
        int endOffset = thisOffset + size;
        for (int i = thisOffset; i < endOffset; ++i) {
            setFloat(i, mapWithIndexFunction.apply(getFloat(i), i));
        }
        return this;
    }

    /**
     * 对向量进行求和
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     * @return
     */
    @Override
    public FloatTensor addInPlace(int thisOffset, FloatTensor that, int thatOffset, int size) {
        return mapWithIndexInPlace(thisOffset, size, (value, index) -> value + that.getFloat(index - thisOffset + thatOffset));
    }

    /**
     * 对向量进行求和
     *
     * @param that
     * @return
     */
    @Override
    public FloatTensor addInPlace(FloatTensor that) {
        return addInPlace(0, that, 0, size());
    }

    /**
     * 对向量进行乘积
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     * @return
     */
    @Override
    public FloatTensor multiplyInPlace(int thisOffset, FloatTensor that, int thatOffset, int size) {
        return mapWithIndexInPlace(thisOffset, size, (value, index) -> value * that.getFloat(index - thisOffset + thatOffset));
    }

    /**
     * 对向量进行乘积
     *
     * @param that
     * @return
     */
    @Override
    public FloatTensor multiplyInPlace(FloatTensor that) {
        return multiplyInPlace(0, that, 0, size());
    }

    /**
     * 对向量进行求商
     *
     * @param thisOffset
     * @param size
     * @param value
     * @return
     */
    @Override
    public FloatTensor divideInPlace(int thisOffset, int size, float value) {
        return mapInPlace(thisOffset, size, f -> f / value);
    }

    /**
     * 填充所有元素为value
     *
     * @param thisOffset 开始位置
     * @param size       填充数量
     * @param value      填充值
     * @return
     */
    @Override
    public FloatTensor fillInPlace(int thisOffset, int size, float value) {
        return mapInPlace(thisOffset, size, unused -> value);
    }

    @Override
    public FloatTensor softmaxInPlace(int thisOffset, int size) {
        // 获取指定长度内值中最大值
        float maxVal = max(thisOffset, size);
        // 值修改为e的 f-最大值
        mapInPlace(thisOffset, size, f -> (float) Math.exp(f - maxVal));
        float sum = sum(thisOffset, size);
        // 值/值的和
        return divideInPlace(thisOffset, size, sum);
    }

    /**
     * 对向量进行a倍乘积+that值
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     * @param a
     * @return
     */
    @Override
    public FloatTensor saxpyInPlace(int thisOffset, FloatTensor that, int thatOffset, int size, float a) {
        // this[thatOffset ... thatOffset + size) = a * that[thatOffset ... thatOffset + size) + this[thisOffset ... thisOffset + size)
        for (int i = 0; i < size; ++i) {
            setFloat(thisOffset + i, a * that.getFloat(thatOffset + i) + this.getFloat(thisOffset + i));
        }
        return this;
    }
}
